import numpy as np


"""
__init__
    bandit: BernoulliBandit类
    counts: 初始化每次的老虎机的拉杆次数, 初始化全部是0
    regret: 当前步数的累计懊悔, 这个值越小则表示每次都选择了最优解
    actions: 记录每一次的动作
    regrets: 记录每一次动作的懊悔
    
update_regret
    计算累积懊悔并保存,k为本次动作选择的拉杆的编号

run_one_step
    子类重写

run
    for 运行 num_steps次数
        选取某根老虎机的拉杆
        这台老虎机的拉杆次数+1
        加入到动作列表
        update_regret计算懊悔并保存到regrets列表

"""
class Solver:
    """ 多臂老虎机算法基本框架 """
    def __init__(self, bandit):
        self.bandit = bandit
        self.counts = np.zeros(self.bandit.K)  # 每根拉杆的尝试次数, 初始化数组里面全部是0
        self.regret = 0.  # 当前步的累积懊悔
        self.actions = []  # 维护一个列表,记录每一步的动作
        self.regrets = []  # 维护一个列表,记录每一步的累积懊悔

    def update_regret(self, k):
        # 计算累积懊悔并保存,k为本次动作选择的拉杆的编号
        self.regret += self.bandit.best_prob - self.bandit.probs[k]
        self.regrets.append(self.regret)

    def run_one_step(self):
        # 返回当前动作选择哪一根拉杆,由每个具体的策略实现
        raise NotImplementedError

    def run(self, num_steps):
        # 运行一定次数,num_steps为总运行次数
        for _ in range(num_steps):
            k = self.run_one_step()
            self.counts[k] += 1
            self.actions.append(k)
            self.update_regret(k)
